Skip to content

fix(streaming): close the SSE event iterators a stream opens - #1

Draft
fancyboi999 wants to merge 7 commits into
mainfrom
fix/close-stream-event-iterator
Draft

fancyboi999 wants to merge 7 commits into
mainfrom
fix/close-stream-event-iterator

Conversation

@fancyboi999

@fancyboi999 fancyboi999 commented Oct 24, 2026 •

Copy link
Copy Markdown
Owner

Implementation and tests for openai#4174, kept in our fork because upstream accepts pull requests only from collaborators.

Changes being requested

Streams now close the iterators they open instead of leaving them to the garbage collector. asyncio closes nested async generators that become unreachable together in concurrent aclose() tasks. One of those tasks then fails with RuntimeError: aclose(): asynchronous generator is already running, logged as Task exception was never retrieved, and that generator's cleanup never runs. We see this in production with Python 3.13 and httpx2 2.13.1. The httpx2-internal part is pydantic/httpx2#1280, fixed by pydantic/httpx2#1281.

  • Stream / AsyncStream
    • _iterator, _iter_events, the decoder iterator and the byte iterator are each closed by the layer that opened it.
    • close() also closes the stream iterator when it is not running.
    • __iter__ / __aiter__ return the stream iterator directly instead of wrapping it in another generator.
    • The public Iterator / AsyncIterator contracts are unchanged.
  • _streaming_cleanup.close_owned / aclose_owned / run_cleanup: one cleanup policy used everywhere.
    • Every resource gets a close attempt, and objects without close / aclose are skipped.
    • Async cleanup is shielded from AnyIO cancellation.
    • An exception already propagating (including one leaving a with block) wins over a cleanup error.
    • A cancellation raised during cleanup is never dropped in favor of an ordinary error.
    • Normal closure (GeneratorExit, [DONE]) does not hide cleanup errors.
  • Chat and Responses stream wrappers and managers close their own iterator and the raw stream.
  • Agent streams
    • Close the parent stream iterator.
    • Pass result errors as the primary failure.
    • Report a terminal cleanup failure after the turn outcome, exactly once.

Known limits, unchanged by this PR and documented in comments or docstrings:

  • close() does not cancel or join a consumer that is iterating concurrently.
  • Shielded cleanup has no deadline.

Additional context & links

Async generators finalized by the GC, from the repro in the issue (local SSE server, Python 3.13.16):

openai / httpx2 read to [DONE] early exit + close()
3.37.0 / 2.16.0 _iter_events, SSEDecoder.aiter_bytes, Response.aiter_bytes, PoolByteStream.__aiter__ AsyncStream.__aiter__
this branch / 2.16.0 PoolByteStream.__aiter__ (httpx2-internal) PoolByteStream.__aiter__
this branch / httpx2#1281 HTTP11ConnectionByteStream.__aiter__ (httpcore2-internal) same

Tests, against main 718aee9 with Python 3.10 and the repo's uv lock:

  • 41 new test cases cover:

    • Generators left to the GC for raw streams and the Chat wrapper.
    • Decoder iterators closed exactly once.
    • Error and cancellation precedence, including caller exceptions and task cancellation at with exit.
    • Agent terminal cleanup.

    All fail on main or on the earlier commits of this branch where the gap existed, and pass at HEAD.

  • Whole suite except tests/api_resources (needs the Steady mock server) and tests/lib/live: 10201 passed on this branch and 10160 on main. The same 200 failures and errors occur on both, from extras not installed here.

  • ruff check and ruff format are clean. mypy reports no issues. pyright 1.1.414 reports 0 errors.

  • An independent model review (seven passes) ended with approve.

Not run: tests/api_resources.

AsyncStream stopped at [DONE] (or was closed early) without closing the
async generators it had opened: _iter_events, the decoder iterator and
the response.aiter_bytes() iterator, plus the __aiter__ wrapper after an
early exit. They stayed suspended until the garbage collector reached
them, where asyncio closes each one in its own aclose() task; a nested
chain closed concurrently fails with "aclose(): asynchronous generator
is already running" and skips its cleanup.

Each layer now closes what it opened: _iter_events closes the decoder
and byte iterators, __stream__ closes the event iterator before the
response, close() closes the stream iterator when it is not running,
and __aiter__ returns the stream iterator instead of wrapping it. The
sync Stream gets the symmetric change.
- Keep the Iterator/AsyncIterator contracts so existing __stream__
  overrides stay valid; close owned iterators through close_owned /
  aclose_owned, which skip objects without close()/aclose().
- Do not let a cleanup error replace the error or cancellation already
  propagating, and shield async cleanup from caller cancellation.
- Create the byte and decoder iterators inside the exception-translation
  block, and loop explicitly in the sync stream so the decoder iterator
  is closed once.
- close() closes the response even when iterator cleanup fails.
- Chat and Responses stream wrappers close their own iterator and the raw
  stream instead of only the response; agent streams close the parent
  stream iterator they consume.
- Tests cover the byte iterator, the chat wrapper, single close of custom
  decoder iterators, and error precedence.
- GeneratorExit from close() or [DONE] is normal closure; it no longer
  counts as a primary failure that silences cleanup errors.
- close_owned/aclose_owned look up and call each closer inside the same
  handler, attempt every resource even after a BaseException, and never
  drop a cancellation raised during cleanup in favor of an ordinary
  error.
- The response joins the same cleanup call as the iterators, so it is
  shielded and follows the same error precedence.
- Comments no longer claim that closing the response ends a concurrent
  consumer.
- Tests check [DONE] cleanup without an explicit close(), cleanup errors
  surfacing on normal closure, and aclose_owned continuing after
  cancellation.
…p errors

- Cleanup helpers let a later control-flow exception (for example
  cancellation) replace an earlier ordinary cleanup failure.
- Agent streams pass the propagating error into cleanup in
  get_final_result(), the event iterator and failed entry, so a decoder
  close failure no longer replaces AgentTurnResultError.
- run_cleanup() covers owners whose async close method is not aclose().
- Tests cover cancellation after an ordinary cleanup error and a failed
  agent turn whose decoder close also fails.
AgentSessionStream closes its connection when it sees the terminal event,
before yielding it. A cleanup failure there was raised before the caller
evaluated the outcome, turning a failed turn or an output parse error
into observation_failed. The failure is now kept and raised only after
the outcome is established (or on the next read when iterating by hand);
a result error still wins.

close() also closes the stream's own iterator when it is not running.
The terminal close marks the stream closed, so the next read stopped
before reaching the deferred cleanup failure and ordinary iteration,
until_done() and close-after-terminal all hid it. Reads on a closed
stream and close() now raise the pending failure once. A result error
from get_final_result() or an exception leaving the with block still
takes precedence and discards it.
…errors

Context-manager exits closed streams without passing the exception that
was leaving the block, so a cleanup failure could replace the caller's
exception or a task cancellation. Stream, AsyncStream, the Chat and
Responses streams and their managers, and the agent session streams now
pass it as the primary failure.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant